# -*- coding:utf-8 -*-#
# @Time:2023/5/31 14:46
# @Author:Adong
# @Software:PyCharm
"""
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
重要文件，谨慎修改！
"""

from torch.utils.data import Dataset
from torchvision import transforms
from PIL import Image
import os

class MyData_V1(Dataset):
    def __init__(self, root_dir, label_dir):
        '''
        数据集的组织架构是：train-[lowSNR_normal,normal,pianci_no,pianci_yes,songdong]-[1.png,2.png,3.png,...]
        :param root_dir:上述train的目录，该目录下包含各类故障的文件夹
        :param label_dir: 数据集的标签
        '''
        self.root_dir = root_dir                                #读取数据集目录,比如训练集或者测试集的目录
        self.label_dir = label_dir                              #读取训练集标签
        self.path = os.path.join(self.root_dir, self.label_dir) #拼接数据集目录和数据集标签，组装标签数据集目录。如r'dataset/CRWU/train/IR0007'
        self.img_name = os.listdir(self.path)                   #读取标签数据集文件名.如’198IR0007.png‘
        self.faulttype = os.listdir(root_dir)                   #读取故障类型
        self.data_transform = transforms.Compose([transforms.ToTensor()])

    def __getitem__(self, idx):                                 #魔术方法，不需要显式调用
        img_name = self.img_name[idx]                           #根据索引idx读取数据集文件名
        img_item_path = os.path.join(self.path, img_name)       #拼接标签数据集目录和标签数据集文件名。如r'dataset/CRWU/train/IR0007/198IR0007.png'
        img = Image.open(img_item_path)                         #打开文件
        label = self.label_dir                                  #记录标签内容
        index = self.faulttype.index(label)                     #把字符串型的标签转换成数字
        # return img,index
        return self.data_transform(img), index      #DataLoader会调用这个函数并接收返回值

    def __len__(self):
        return len(self.img_name)

class MyData_V2(Dataset):
    '''
    转为自编码器涉及的dataset
    '''
    def __init__(self,root_dir):
        self.root_dir = root_dir
        self.img_name = os.listdir(self.root_dir)
        self.data_transform = transforms.Compose([transforms.ToTensor()])

    def __getitem__(self, idx):
        img_name = self.img_name[idx]
        img_item_path = os.path.join(self.root_dir,img_name)
        img = Image.open(img_item_path)
        source= self.data_transform(img)
        target = self.data_transform(img)
        return source,target
    def __len__(self):
        return len(self.img_name)


def makeDataset(mode,dataset_verision):
    '''
    制作有标注的数据集,MyData_V1
    :param mode:
    :param verision:
    :return:
    '''
    if mode == 'train':
        root_dir = r'./data/' + dataset_verision + '/train/'
    elif mode == 'test':
        root_dir = r'./data/' + dataset_verision + '/test/'
    elif mode == 'veri':
        root_dir = r'./data/' + dataset_verision + '/veri/'
    else:
        print('what the fuck U input?')
    faulttype = os.listdir(root_dir)
    dataset = MyData_V1(root_dir, faulttype[0])
    for t in faulttype[1:]:
        dataset = dataset + MyData_V1(root_dir, t)
    return dataset


if __name__ == '__main__':
    trainset = makeDataset('train','wavV3_to_Mel')
    # testset = makeDataset('test','wav_data_V4')
    # veriset = makeDataset('veri','wav_data_V4')
    # normal = MyData_V2(r'./data/wav_to_gray/train/normal')